from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import array, random, linspace, pi, ravel, cos, sin, empty, mgrid, arange
from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface
import csv
from scipy.special import sph_harm

engine = Engine()
engine.start()


class Form(HasTraits):
    scene = Instance(MlabSceneModel, args={'engine': engine})
    file = File('temperature.axdt')

    button = Button('显示温度场')

    def traits_init(self):
        self.mymlab = self.scene.mlab
    def _button_fired(self):
        with open(self.file) as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        self.v = v
        self.f = f
        self.t = t
        mymlab = self.mymlab
        mymlab.figure(self.scene.mayavi_scene)
        mymlab.clf()
        surf = mymlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0))
        mymlab.colorbar()

    buttonIR = Button('计算红外辐射')

    def _buttonIR_fired(self):
        v = self.v
        f = self.f
        t = self.t
        thetas = np.linspace(-40, 40, 41)
        phis = np.linspace(0, 360, 181)
        ir = np.zeros((thetas.shape[0], phis.shape[0]))
        v0 = v[f[:, 0], :]
        v1 = v[f[:, 1], :]
        v2 = v[f[:, 2], :]
        tf = (t[f[:, 0]] + t[f[:, 1]] + t[f[:, 2]]) / 3
        n = np.cross(v1 - v0, v2 - v0)
        a = np.sqrt((n * n).sum(1)) / 2
        n /= 2 * a[:, np.newaxis]
        for tt, theta in enumerate(thetas):
            costheta = np.cos(theta * np.pi / 180)
            sintheta = np.sin(theta * np.pi / 180)
            print(tt)
            for pp, phi in enumerate(phis):
                cosphi = np.cos(phi * np.pi / 180)
                sinphi = np.sin(phi * np.pi / 180)
                s = np.array([costheta * cosphi, costheta * sinphi, sintheta])
                ns = n @ s
                light = ns < -1e-6
                ir[tt, pp] = np.sum(5.67e-8 * tf[light] ** 4 / np.pi * (-ns[light]) * a[light])
        self.ir = ir
        self.thetas = thetas
        self.phis = phis
        print(ir.max())

    buttonShowIR = Button('显示红外辐射')

    def _buttonShowIR_fired(self):
        mymlab = self.mymlab
        thetas, phis = np.meshgrid(self.thetas, self.phis, indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir = self.ir * 5
        mymlab.mesh(ir * x, ir * y, ir * z, opacity=.5)
        mymlab.colorbar()

    buttonShowIRSphere = Button('显示红外辐射球面')

    def _buttonShowIRSphere_fired(self):
        mymlab = self.mymlab
        thetas, phis = np.meshgrid(self.thetas, self.phis, indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir = self.ir * 5
        mymlab.mesh(x, y, z, scalars=ir, opacity=.5)
        mymlab.colorbar()

    buttonClear = Button('清除视景')

    def _buttonClear_fired(self):
        mymlab = self.mymlab
        mymlab.clf()
        mymlab.draw()

    buttonShowIRScene = Button('测试')

    def _buttonShowIRScene_fired(self):
        mymlab = self.mymlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # mymlab.figure()
        surf = mymlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        mymlab.draw()

    buttonTest1 = Button('测试1')

    def _buttonTest1_fired(self):
        mymlab = self.mymlab

        def polydata():
            # The numpy array data.
            points = np.array([[0, -0.5, 0], [1.5, 0, 0], [0, 1, 0], [0, 0, 0.5], [-1, -1.5, 0.1], [0, -1, 0.5], [-1, -0.5, 0], [1, 0.8, 0]], 'f')
            triangles = np.array([[0, 1, 3], [1, 2, 3], [1, 0, 5], [2, 3, 4], [3, 0, 4], [0, 5, 4], [2, 4, 6], [2, 1, 7]])
            scalars = np.random.random(points.shape)
            # The TVTK dataset.
            mesh = tvtk.PolyData(points=points, polys=triangles)
            mesh.point_data.scalars = scalars
            mesh.point_data.scalars.name = 'scalars'
            return mesh

        def view(dataset):
            surf = mymlab.pipeline.surface(dataset, opacity=0.1)
            mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(polydata())

    buttonTest2 = Button('测试2')

    def _buttonTest2_fired(self):
        mymlab = self.mymlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # scalars = np.array([[0, 1, 0]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.point_data.scalars = scalars
        # mesh.cell_data.scalars = scalars
        surf = mymlab.pipeline.surface(mesh)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest3 = Button('测试3')

    def _buttonTest3_fired(self):
        mymlab = self.mymlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([[0, .5, 1]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.cell_data.scalars = scalars
        surf = mymlab.pipeline.surface(mesh)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest4 = Button('测试4')

    def _buttonTest4_fired(self):
        mymlab = self.mymlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        surf = mymlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest5 = Button('测试5，默认红外文件')

    def _buttonTest5_fired(self):

        with open('temperature.axdt') as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        mymlab = self.mymlab
        surf = mymlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mymlab.colorbar()
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        pass

    v = None
    f = None
    t = None

    view = View(
        UItem('scene', editor=SceneEditor(scene_class=MayaviScene)),
        HGroup(
            UItem('button'),
            UItem('buttonIR'),
            UItem('buttonClear'),
            UItem('buttonShowIR'),
            UItem('buttonShowIRScene'),
            UItem('buttonShowIRSphere'),
            UItem('buttonTest1'),
            UItem('buttonTest2'),
            UItem('buttonTest3'),
            UItem('buttonTest4'),
            UItem('buttonTest5'),
            UItem('file'),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )


from PyQt5.QtWidgets import *


class FormWidget(QWidget):
    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)


formWidget = FormWidget()
if __name__ == '__main__':
    import os
    import sys

    app = QApplication(sys.argv)
    formWidget.show()
    sys.exit(app.exec_())
